#pragma once

#include <WS2tcpip.h>
#include <functional>
#include <map>

#include "conn.h"

class TcpClient {
public:

    TcpClient() {
    }

    void Connect(const char *addr, int port, int family = IpFamily::v4) {
        WSAData wsData;
        if (WSAStartup(MAKEWORD(2, 2), &wsData) != 0) {
            throw std::runtime_error("WSAStartup failed: " + std::to_string(WSAGetLastError()));
        }

        struct sockaddr_in targetAddr;
        targetAddr.sin_family = family;
        inet_pton(family, addr, &targetAddr.sin_addr);
        targetAddr.sin_port = htons(port);

        so_ = new _Socket();
        SOCKET s = socket(family, SOCK_STREAM, 0);
        if (s == Err::InvalidSocket) {
            throw std::runtime_error("Could not create socket: " + std::to_string(WSAGetLastError()));
        }
        so_->s = s;

        if (connect(so_->s, (struct sockaddr *) &targetAddr, sizeof(targetAddr)) < 0) {
            throw std::runtime_error("Connection error: " + std::to_string(WSAGetLastError()));
        }

        u_long nonBlock = 1;
        if (!SetNonBlock(so_->s, &nonBlock)) {
            throw std::runtime_error("ioctlsocket() failed with error: " + std::to_string(WSAGetLastError()));
        }

        conn_ = new TcpConn(so_);
    }

    bool Run() {

        FD_ZERO(&readSet_);
        FD_ZERO(&writeSet_);
        FD_SET(conn_->Socket(), &readSet_);
        FD_SET(conn_->Socket(), &writeSet_);

        // Select
        int ret = select(0, &readSet_, &writeSet_, nullptr, nullptr);
        if (ret == Err::SocketError) {
            throw std::runtime_error("Select Error: " + std::to_string(WSAGetLastError()));
        }

        if (FD_ISSET(conn_->Socket(), &readSet_)) {
            static char *BUF = new char[MAX_PACK_SIZE];
            int16_t proto = 0;
            int size = 0;
            if ((ret = conn_->Read(&proto, BUF, &size)) == 0) {
                auto h = handlers_.find(proto);
                if (h == handlers_.end()) {
                    std::cout << "Error: " << "Cannot find the handler of proto " << proto << std::endl;
                }
                conn_->SetProto(proto);
                h->second(conn_, BUF, size);
            } else {
                std::cout << GetErrorMsg(ret) << std::endl;
                return false;
            }
        }
        return true;
    }

    int Read(char *data, int len) {
        return conn_->Read(data, len);
    }

    int Write(const char *data, int len) {
        return conn_->Write(data, len);
    }

    TcpConn *Conn() {
        return conn_;
    }

    template<class T>
    void Register(int16_t proto, HandlerFunc<T> handler) {
        const auto &iter = handlers_.find(proto);
        if (iter == handlers_.end()) {
            T *instance = new T();
            handlers_[proto] = [instance, handler](TcpConn *conn, const char *data, int size) {
                instance->ParseFromArray(data, size);
                std::string json;
                if(google::protobuf::util::MessageToJsonString(*instance, &json) == google::protobuf::util::Status::OK) {
                    std::cout << "recv " << typeid(*instance).name() << " " << json << std::endl;
                }
                handler(conn, *instance);
            };
        } else {
            throw std::logic_error("duplicated register " + std::to_string(proto));
        }
    }

    ~TcpClient() {
        delete conn_;
        Close(so_);
    }

private:
    TcpConn *conn_;
    _Socket *so_;
    std::map<int16_t, std::function<void(TcpConn *, const char *, int)>> handlers_;
    fd_set readSet_;
    fd_set writeSet_;
};
